# -*- coding: utf-8 -*-
import pandas as pd
import xgboost as xgb
import talib
import efinance as ef
import sys

# --- 1. 从命令行获取参数 ---
if len(sys.argv) < 3:
    print("错误：参数不足。")
    print("用法: python3 05_predict_next.py <股票代码> <frequency>")
    sys.exit(1)

stock_code = sys.argv[1]
frequency = sys.argv[2]

# --- 2. 定义模型文件名 ---
model_path = f'output/models/xgb_model_{stock_code}_{frequency}.json'

# --- 3. 加载已训练好的模型 ---
try:
    print(f"加载已保存的 XGBoost 模型: {model_path}...")
    model = xgb.XGBRegressor()
    model.load_model(model_path)
    print("模型加载成功。")
except xgb.core.XGBoostError:
    print(f"错误: 找不到模型文件 {model_path}。")
    print(f"请先运行 'python3 run_technical_analysis.py train {stock_code} -f {frequency}' 来训练模型。")
    sys.exit(1)

# --- 4. 根据 frequency 设置 klt 参数 ---
klt = 0
if frequency == 'daily':
    klt = 101
elif frequency == 'hourly':
    klt = 60
else:
    print(f"错误：不支持的时间周期 '{frequency}'。")
    sys.exit(1)

# --- 5. 获取最新的股票数据 ---
print(f"开始获取 {stock_code} 最新的 {frequency} 数据...")
df_latest = ef.stock.get_quote_history(stock_code, klt=klt, end='20300101')
if df_latest.empty:
    print(f"错误：未能获取到 {stock_code} 的最新数据。")
    sys.exit(1)
print("最新数据获取成功。")

# --- 6. 准备数据与计算指标 ---
stock_name = df_latest['股票名称'].iloc[0]
date_col = '日期' if '日期' in df_latest.columns else '时间'
df_latest.rename(columns={
    date_col: 'date', '开盘': 'open', '收盘': 'close', '最高': 'high',
    '最低': 'low', '成交量': 'volume'
}, inplace=True)
df_latest['date'] = pd.to_datetime(df_latest['date'])
df_latest.set_index('date', inplace=True)

print("开始为最新数据计算技术指标...")
df_latest['sma5'] = talib.SMA(df_latest['close'], timeperiod=5)
df_latest['sma20'] = talib.SMA(df_latest['close'], timeperiod=20)
df_latest['ema12'] = talib.EMA(df_latest['close'], timeperiod=12)
df_latest['ema26'] = talib.EMA(df_latest['close'], timeperiod=26)
df_latest['macd'], df_latest['macdsignal'], df_latest['macdhist'] = talib.MACD(df_latest['close'], fastperiod=12, slowperiod=26, signalperiod=9)
df_latest['rsi'] = talib.RSI(df_latest['close'], timeperiod=14)
df_latest['upperband'], df_latest['middleband'], df_latest['lowerband'] = talib.BBANDS(df_latest['close'], timeperiod=20, nbdevup=2, nbdevdn=2, matype=0)
print("技术指标计算完成。")

# --- 7. 准备用于预测的最后一行数据 ---
feature_columns = [
    'volume', 'sma5', 'sma20', 'ema12', 'ema26', 'macd', 'macdsignal',
    'macdhist', 'rsi', 'upperband', 'middleband', 'lowerband'
]
latest_features = df_latest[feature_columns].iloc[-1:]

if latest_features.isnull().values.any():
    print("\n警告：最新的特征数据中包含缺失值 (NaN)。预测可能不准确。")
else:
    # --- 8. 进行预测 ---
    print("\n开始进行预测...")
    prediction = model.predict(latest_features)
    predicted_price = prediction[0]
    latest_close_price = df_latest['close'].iloc[-1]
    
    predict_unit = "明天" if frequency == 'daily' else "下一个小时"
    
    print("\n--- 预测结果 ---")
    print(f"公司名称: {stock_name} ({stock_code})")
    print(f"时间周期: {'日线' if frequency == 'daily' else '60分钟线'}")
    print(f"当前周期的收盘价是: {latest_close_price:.2f}")
    print(f"模型预测{predict_unit}的收盘价是: {predicted_price:.2f}")

    change = predicted_price - latest_close_price
    change_percent = (change / latest_close_price) * 100

    print(f"预测涨跌幅: {change_percent:.2f}%")
    if change > 0:
        print("预测结果: 看涨 📈")
    else:
        print("预测结果: 看跌 📉")